# encoding utf-8

from __future__ import annotations

from typing import List, Tuple
import pandas as pd

from core.database import db_manager, DatabaseManager
from dao.models import MarginData


class MarginDataDAO:
    """融资融券数据访问对象"""

    def __init__(self, db: DatabaseManager | None = None):
        self.db = db or db_manager

    def get_margin_data(self, exchange: str, start_date: str, end_date: str) -> pd.DataFrame:
        """从数据库获取融资融券数据"""
        if exchange not in ['sz', 'sh']:
            return pd.DataFrame()

        def normalize_date(date_str: str) -> str:
            if len(date_str) == 8 and date_str.isdigit():
                return f"{date_str[:4]}-{date_str[4:6]}-{date_str[6:8]}"
            return date_str

        start_date = normalize_date(start_date)
        end_date = normalize_date(end_date)

        session = self.db.get_session()
        try:
            ex_name = 'SZSE' if exchange == 'sz' else 'SSE'
            rows = (
                session.query(MarginData)
                .filter(MarginData.exchange == exchange)
                .filter(MarginData.date >= start_date)
                .filter(MarginData.date <= end_date)
                .order_by(MarginData.date)
                .all()
            )
            if not rows:
                return pd.DataFrame()
            df = pd.DataFrame([
                {
                    'date': r.date,
                    'financing_balance': r.financing_balance,
                    'financing_buy_amount': r.financing_buy_amount,
                    'securities_lending_volume': r.securities_lending_volume,
                    'securities_lending_balance': r.securities_lending_balance,
                    'securities_lending_sell_volume': r.securities_lending_sell_volume,
                    'total_margin_balance': r.total_margin_balance,
                }
                for r in rows
            ])
            df['exchange_name'] = ex_name
            return df
        finally:
            session.close()

    def save_margin_data(self, exchange: str, data: pd.DataFrame):
        """保存数据到数据库"""
        if exchange not in ['sz', 'sh'] or data.empty:
            return
        session = self.db.get_session()
        try:
            for _, row in data.iterrows():
                key_date = str(row['date'])
                item = session.get(MarginData, {"date": key_date, "exchange": exchange})
                if not item:
                    item = MarginData(date=key_date, exchange=exchange)
                    session.add(item)
                item.financing_balance = float(row['financing_balance']) if row['financing_balance'] is not None else None
                item.financing_buy_amount = float(row['financing_buy_amount']) if row['financing_buy_amount'] is not None else None
                item.securities_lending_volume = float(row['securities_lending_volume']) if row['securities_lending_volume'] is not None else None
                item.securities_lending_balance = float(row['securities_lending_balance']) if row['securities_lending_balance'] is not None else None
                item.securities_lending_sell_volume = float(row['securities_lending_sell_volume']) if row['securities_lending_sell_volume'] is not None else None
                item.total_margin_balance = float(row['total_margin_balance']) if row['total_margin_balance'] is not None else None
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()


# 创建DAO实例
margin_data_dao = MarginDataDAO()

